Skip to content

Return all embeddings for batched OpenAI inputs - #71989

Merged
kaxil merged 1 commit into
apache:mainfrom
YAshhh29:fix-openai-batched-embeddings
Sep 17, 2026
Merged

kaxil merged 1 commit into
apache:mainfrom
YAshhh29:fix-openai-batched-embeddings

Conversation

@YAshhh29

@YAshhh29 YAshhh29 commented Aug 23, 2026 •

Copy link
Copy Markdown
Contributor

Why

The OpenAI Embeddings API accepts batches, but OpenAIHook.create_embeddings() returned only response.data[0]. Every successful batch result after the first item was silently discarded.

What changed

  • Preserve the existing single-input return shape.
  • Return every batch embedding ordered by the response index.
  • Update the operator return type and user documentation.
  • Add regressions for string and token-array batches.

Tests

  • Focused batch regression test against this branch: 2 passed.
  • The same regression test against the unmodified implementation: 2 failed.
  • ruff format and ruff check --fix: passed.
  • git diff --check: passed.

The full provider suite was not run locally because Airflow does not support native Windows; CI provides the Linux suite result.


Was generative AI tooling used to co-author this PR?
  • Yes - GitHub Copilot

Generated-by: GitHub Copilot following the guidelines


Drafted-by: GitHub Copilot and reviewed by @YAshhh29

The SDK returns one indexed embedding per batch item, but the hook exposed only the first result. Batch callers therefore received a plausible-looking response with no indication that every later input had been discarded.
@kaxil
kaxil merged commit d9a5f77 into apache:main Sep 17, 2026
79 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants